package tree;

import java.util.Comparator;

public class AVLTree<E> extends BST<E> {

	public AVLTree() {
		this(null);
	}

	public AVLTree(Comparator<E> comparator) {
		super(comparator);
	}

	@Override
	protected Node<E> createNode(E element, Node<E> parent) {
		return new AVLNode<>(element, parent);
	}

	@Override
	protected void afterAdd(Node<E> node) {// 增加叶子节点以后执行
		while ((node = node.parent) != null) {// 新增加的节点的父节点不为空一直while循环
			if (isBalanced(node)) {// 如果叶子结点的高度平衡 左右差<=1
				updateHeight(node); // 更新节点高度，因为是二叉搜索树，添加后有个父节点必然为2；
			} else {
				reBalanced2(node);
				break;
			}
		}
	}

	private void updateHeight(Node<E> node) {
		((AVLNode<E>) node).updateHeight();
	}

	private void reBalanced(Node<E> grand) {
		Node<E> parent = ((AVLNode<E>) grand).tallerNode();
		Node<E> child = ((AVLNode<E>) parent).tallerNode();
		if (parent.isLeftChild()) { // L
			if (child.isLeftChild()) { // LL
				rotateRight(grand);
			} else { // LR
				rotateLeft(parent);
				rotateRight(grand);
			}
		} else { // R
			if (child.isLeftChild()) { // RL
				rotateRight(parent);
				rotateLeft(grand);
			} else { // RR
				rotateLeft(grand);
			}
		}
	}

	private void reBalanced2(Node<E> grand) {
		Node<E> parent = ((AVLNode<E>) grand).tallerNode();
		Node<E> child = ((AVLNode<E>) parent).tallerNode();
		if (parent.isLeftChild()) { // L
			if (child.isLeftChild()) { // LL
				rotate(grand, child.left, child, child.right, parent, parent.right, grand, grand.right);
			} else { // LR
				rotate(grand, parent.left, parent, child.left, child, child.right, grand, grand.right);
			}
		} else { // R
			if (child.isLeftChild()) { // RL
				rotate(grand, grand.left, grand, child.left, child, child.right, parent, parent.right);
			} else { // RR
				rotate(grand, grand.left, grand, parent.left, parent, child.left, child, child.right);
			}
		}
	}

	private void rotateRight(Node<E> grand) {
		Node<E> parent = grand.left;
		Node<E> child = parent.right;

		parent.right = grand;
		grand.left = child;

		afterRotate(grand, parent, child);
	}

	private void rotateLeft(Node<E> grand) {
		Node<E> parent = grand.right;
		Node<E> child = parent.left;
		grand.right = child;
		parent.left = grand;

		afterRotate(grand, parent, child);
	}

	private void afterRotate(Node<E> grand, Node<E> parent, Node<E> child) {
		parent.parent = grand.parent;
		if (grand.isLeftChild()) {
			grand.parent.left = parent;
		} else if (grand.isRightChild()) {
			grand.parent.right = parent;
		} else {
			root = parent;
		}
		if (child != null) {
			child.parent = grand;
		}
		grand.parent = parent;
		updateHeight(grand);
		updateHeight(parent);
	}

	private void rotate(Node<E> r, Node<E> a, Node<E> b, Node<E> c, Node<E> d, Node<E> e, Node<E> f, Node<E> g) {
		d.parent = r.parent;
		if (r.isLeftChild()) {
			r.parent.left = d;
		} else if (r.isRightChild()) {
			r.parent.right = d;
		} else {
			root = d;
		}
		b.left = a;
		b.right = c;
		if (a != null) {
			a.parent = b;
		}
		if (c != null) {
			c.parent = b;
		}
		updateHeight(b);
		f.left = e;
		f.right = g;
		if (e != null) {
			e.parent = f;
		}
		if (g != null) {
			g.parent = f;
		}
		updateHeight(f);
		d.left = b;
		d.right = f;
		b.parent = d;
		f.parent = d;
		updateHeight(d);
	}

	@Override
	protected void afterRemove(Node<E> node) {
		while ((node = node.parent) != null) {// 新增加的节点的父节点不为空一直while循环
			if (isBalanced(node)) {// 如果叶子结点的高度平衡 左右差<=1
				updateHeight(node); // 更新节点高度，因为是二叉搜索树，添加后有个父节点必然为2；
			} else {
				reBalanced2(node);

			}
		}
	}

	private boolean isBalanced(Node<E> node) {
		return Math.abs(((AVLNode<E>) node).balanceFactor()) <= 1;
	}

//内部类 继承
	protected static class AVLNode<E> extends Node<E> {
		int height = 1;

		public AVLNode(E element, Node<E> parent) {
			super(element, parent);
		}

		public int balanceFactor() {
			int leftHeight = left == null ? 0 : ((AVLNode<E>) left).height;
			int rightHeight = right == null ? 0 : ((AVLNode<E>) right).height;
			return leftHeight - rightHeight;
		}

		public void updateHeight() {
			int leftHeight = left == null ? 0 : ((AVLNode<E>) left).height;
			int rightHeight = right == null ? 0 : ((AVLNode<E>) right).height;
			height = Math.max(leftHeight, rightHeight) + 1;
		}

		public Node<E> tallerNode() {
			int leftHeight = left == null ? 0 : ((AVLNode<E>) left).height;
			int rightHeight = right == null ? 0 : ((AVLNode<E>) right).height;
			return leftHeight > rightHeight ? left : right;
		}

		@Override
		public String toString() {
			String parentString = "null";
			if (parent != null) {
				parentString = parent.element.toString();
			}
			return element + "_p(" + parentString + ")_h(" + height + ")";
		}
	}

}
